import numpy as np

# 添加椒盐噪声的函数
def add_salt_and_pepper_noise(img, salt_prob=0.0005):
    img = img.astype(np.float32)
    h, w, c = img.shape
    num_pixels = int(h * w * salt_prob)

    # 盐：白点
    coords = [np.random.randint(0, i - 1, num_pixels) for i in (h, w)]
    img[coords[0], coords[1], :] = 255.0

    # 椒：黑点
    coords = [np.random.randint(0, i - 1, num_pixels) for i in (h, w)]
    img[coords[0], coords[1], :] = 0.0
    return img

